feat(rust/signal): idempotent maintain_keys for the pre-key lifecycle - #267
Conversation
…ycle Adds `signal::maintain::maintain_keys` — a health-gated, idempotent reconcile of an agent's published pre-key bundle with the relay. It publishes only when needed: a no-op when the store already holds a signed pre-key and the relay's one-time pool is healthy, a (re)publish on first boot, a wiped store, or a low/depleted pool. Storing private material locally before upload, and skipping the upload when the store already backs a healthy set, is the anti-orphan invariant — blindly republishing on every start leaves the relay advertising one-time pre-keys a rotated/wiped store can no longer back, so peers' X3DH fails with a MAC error. Composes existing pieces (KeysApi health/rotate/upload + signal::keys generation + SessionStore), so callers stop hand-rolling the lifecycle. Adds wiremock coverage for the idempotent no-op and the depleted-pool refill paths. Co-Authored-By: Claude <noreply@anthropic.com>
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds a public Rust API for idempotent Signal pre-key maintenance. It reconciles relay health with locally backed keys, conditionally rotates signed pre-keys, replenishes one-time pre-keys, and reports maintenance results. Integration tests cover healthy no-ops, targeted repairs, pool refilling, and first-boot publication. ChangesSignal key maintenance
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant maintain_keys
participant KeysApi
participant Signer
participant SessionStore
maintain_keys->>KeysApi: check relay health
maintain_keys->>SessionStore: check signed-key backing
maintain_keys->>Signer: generate signed and one-time pre-keys
maintain_keys->>SessionStore: store private key material
maintain_keys->>KeysApi: rotate signed key and upload public pre-keys
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f8e9fc22f0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if store.active_signed_pre_key().await.is_ok() { | ||
| if let Ok(health) = keys.health(agent_id).await { | ||
| if !health.low_one_time_pre_keys && health.one_time_pre_key_count > 0 { |
There was a problem hiding this comment.
When the local store has any active signed pre-key but the relay is advertising a different signedPreKeyKeyId (for example after restoring a stale/partial store or after a failed rotation left the local active id ahead of the server), this branch returns healthy as long as the one-time pool count is fine. The next inbound PREKEY_BUNDLE carries the relay's signed pre-key id, and SignalSession::process_pre_key_message looks up that exact id; if it is absent locally, first-contact decrypt fails instead of maintenance repairing the bundle. Use health.signed_pre_key_key_id to confirm store.signed_pre_key(id) exists before taking the no-op path.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sdk/rust/src/signal/maintain.rs`:
- Around line 83-133: Decouple signed-key and one-time-key maintenance in the
reconciliation flow. Use the active signed pre-key status to rotate and store a
new signed key only when that key is missing or invalid, while independently
checking the relay’s one-time pre-key health and generating/uploading a batch
only when required. Preserve the existing healthy no-op behavior and ensure
one-time-key depletion does not trigger unnecessary signed-key rotation; keep
the existing rotate-failure handling scope unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b9e9801e-1f09-427e-aa54-d0889b6f8387
📒 Files selected for processing (3)
sdk/rust/src/signal/maintain.rssdk/rust/src/signal/mod.rssdk/rust/tests/signal_maintain.rs
Addresses review on #267. The no-op gate previously required BOTH "store has an active signed pre-key" AND "relay pool healthy", so a merely-depleted one-time pool fell through to a full republish that needlessly rotated a perfectly good signed pre-key. Worse, it never checked WHICH signed pre-key the relay advertises: `active_signed_pre_key().is_ok()` is true even when the store's active key is not the one the relay serves, so an inbound PREKEY_BUNDLE naming the relay's id would fail to decrypt and maintenance would never repair it. The two halves are now maintained independently: - Signed pre-key: rotated only when the store cannot answer `KeyHealth::signed_pre_key_key_id` — the exact id an inbound PREKEY_BUNDLE names. A still-backed key is left untouched even when the pool needs a refill. - One-time pre-keys: topped up only on the relay's own low/empty report, without disturbing a healthy signed pre-key. Tests now pin one half healthy and the other degraded to prove neither drags the other into a needless republish, plus the first-boot both-halves case. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
sdk/rust/tests/signal_maintain.rs (1)
27-39: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMatch the health endpoint explicitly.
This mock returns health for any
GET, so an incorrect health path can still pass all scenarios. Match/keys/x/healthto preserve the relay-route contract.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sdk/rust/tests/signal_maintain.rs` around lines 27 - 39, Restrict the GET mock in the signal maintenance test to the exact `/keys/x/health` path by adding the appropriate path matcher alongside method("GET"). Leave the PUT mock unchanged so incorrect health routes no longer satisfy the test.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sdk/rust/tests/signal_maintain.rs`:
- Around line 121-144: Update the test
tops_up_only_the_pool_when_depleted_leaving_a_good_signed_key_alone to use a
nonzero low-pool response, such as relay(4, json!("spk_known")), while
preserving the existing signed-key and pool-only upload assertions. Ensure the
test specifically validates behavior driven by the lowOneTimePreKeys flag rather
than oneTimePreKeyCount <= 0.
---
Nitpick comments:
In `@sdk/rust/tests/signal_maintain.rs`:
- Around line 27-39: Restrict the GET mock in the signal maintenance test to the
exact `/keys/x/health` path by adding the appropriate path matcher alongside
method("GET"). Leave the PUT mock unchanged so incorrect health routes no longer
satisfy the test.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8f6a4199-82eb-4ecd-8a1b-e10662790739
📒 Files selected for processing (2)
sdk/rust/src/signal/maintain.rssdk/rust/tests/signal_maintain.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- sdk/rust/src/signal/maintain.rs
Addresses review on #267. - The GET mock matched any path, so a wrong health route would still satisfy every scenario. Pin it to `^/keys/[^/]+/health$` (the id segment is the agent's base58 key, not a literal). - `tops_up_only_the_pool_when_low` used a zero count, where BOTH disjuncts of `low_one_time_pre_keys || one_time_pre_key_count <= 0` are true, so the low-flag branch was never isolated. Use a positive-but-low count (4) to pin that branch specifically; the empty-pool branch stays covered by the first-boot case. Co-Authored-By: Claude <noreply@anthropic.com>
1 participant
What
Adds
signal::maintain::maintain_keysto the Rust SDK — a health-gated, idempotent reconcile of an agent's published pre-key bundle with the relay.It publishes only when the relay actually needs it:
Why
The current pattern (every consumer hand-rolls "publish on boot") re-mints and re-uploads a fresh bundle on every start. After a rotated or wiped local store that orphans the keys the relay still serves — the relay hands peers a bundle the agent can no longer back, so their X3DH fails with a MAC error and the first message is silently dropped.
Storing private material before upload, and skipping the upload entirely when the store already backs a healthy set, is the anti-orphan invariant. Centralizing it in the SDK (like libsignal's refill-when-low) means every agent gets the correct behavior instead of each reimplementing it.
Where it lives
signal/(which already owns key generation + theSessionStore), notapi/keys.rs— that module is intentionally pure REST and "does not generate or validate any Signal keys itself."maintain_keyscomposes the existingKeysApi(health/rotate/upload) +signal::keysgeneration +SessionStore.Tests
tests/signal_maintain.rs(wiremock): the idempotent no-op path (exactly one upload across two healthy calls) and the depleted-pool refill path. Full Rust SDK suite +cargo fmt --check+cargo clippy -D warningsgreen.Not included
maintain_keysprevents new orphans; garbage-collecting already-orphaned one-time keys off the relay needs a relay-sideDELETE /keys/:id/prekeysendpoint (separate, server-side) that doesn't exist yet.Summary by CodeRabbit